
import pandas as pd
#import numpy as np
from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LinearRegression
import os

# 缺失值处理记录
fill_method_records = []

def fill_missing_values(df, country, indicator, method):
    """根据不同的方法填充缺失值"""
    if method == '1':  # 均值填充
        fill_value = df[indicator].mean()
        df[indicator].fillna(fill_value, inplace=True)
    elif method == '2':  # 前向填充
        df[indicator].fillna(method='ffill', inplace=True)
    elif method == '3':  # 后向填充
        df[indicator].fillna(method='bfill', inplace=True)
    elif method == '4':  # 前后向填充
        if pd.isna(df[indicator].iloc[0]):
            df[indicator].fillna(method='bfill', inplace=True)
        df[indicator].fillna(method='ffill', inplace=True)
    elif method == '5':  # 线性插值
        if pd.isna(df[indicator].iloc[0]):
            df[indicator].fillna(method='bfill', inplace=True)
        df[indicator].interpolate(method='linear', inplace=True)
    elif method == '6':  # 多项式插值（三阶）
        if df[indicator].isna().iloc[0]:
            df[indicator].fillna(method='bfill', inplace=True)
            if df[indicator].isna().iloc[0]:
                df[indicator].fillna(method='ffill', inplace=True)

        df_non_nan = df.dropna(subset=[indicator])
        X = df_non_nan[['Year']]
        y = df_non_nan[indicator]
        poly = PolynomialFeatures(degree=3)
        X_poly = poly.fit_transform(X)
        poly_reg = LinearRegression()
        poly_reg.fit(X_poly, y)

        for idx in df.index:
            if pd.isna(df.loc[idx, indicator]):
                X_pred = poly.transform([[df.loc[idx, 'Year']]])
                pred = poly_reg.predict(X_pred)
                df.loc[idx, indicator] = max(pred[0], 0)
    return df

def process_missing_values(df, index_list, country_list, output_root):
    """对指定的指标和国家进行缺失值处理"""
    global fill_method_records
    for indicator in index_list:
        record_data = []
        for i, country in enumerate(country_list, start=1):
            country_df = df[df['国名Ch'] == country][['国名Ch', indicator, 'Year']].copy()
            stats = country_df[indicator].describe()
            missing_percent = country_df[indicator].isna().mean() * 100

            print(f"\n当前指标数据（{indicator} - {country}）：")
            print(country_df)
            print(f"当前指标描述统计结果：\n{stats}")
            print(f"当前指标缺失值百分比：{missing_percent:.2f}%")

            user_input = input(f"是否需要对指标 {indicator}、国家 {country} 进行缺失值处理？(Y/N): ").strip().upper()
            if user_input == 'N':
                continue

            while True:
                method = input("请选择缺失值处理方法：\n 1. 均值填充\n 2. 前向填充\n 3. 后向填充\n 4. 前后向填充\n 5. 线性插值\n 6. 多项式插值\n请输入: ").strip()
                if method in ['1', '2', '3', '4', '5', '6']:
                    break
                else:
                    print("无效输入，请重新输入 1-6 之间的数字。")

            method_dict = {
                '1': '均值填充',
                '2': '前向填充',
                '3': '后向填充',
                '4': '前后向填充',
                '5': '线性插值',
                '6': '多项式插值'
            }
            record_data.append([i, country, method_dict[method]])

            df.loc[df['国名Ch'] == country, indicator] = fill_missing_values(
                df.loc[df['国名Ch'] == country], country, indicator, method)[indicator]

            print(f"处理后数据（{indicator} - {country}）：")
            print(df[df['国名Ch'] == country][['国名Ch', indicator, 'Year']])

        # 保存记录和结果
        fill_record_dir = os.path.join(output_root, "fill_records")
        fill_result_dir = os.path.join(output_root, "filled_result")
        os.makedirs(fill_record_dir, exist_ok=True)
        os.makedirs(fill_result_dir, exist_ok=True)

        records_df = pd.DataFrame(record_data, columns=['序号', '国名Ch', '填充方法'])
        records_df.to_csv(os.path.join(fill_record_dir, f"{indicator}_fill_record.csv"), index=False, encoding='utf-8-sig')
        df[['序号', '国名Ch', indicator, 'Year']].to_csv(os.path.join(fill_result_dir, f"{indicator}_filled.csv"), index=False, encoding='utf-8-sig')
